"""
Description: A simple teleoperation example using keyboard input to control a robot in simulation.

Author: Jin Liu
Date: 2025/06/18
"""

from ezgl import tb, Matrix4x4
from ezrobot.sim import RobotClient
from ezrobot.utils import GraphRunner
from ezrobot.motion import OnlineMotion


if __name__ == '__main__':
    client = RobotClient()
    client.clear_items()
    client.toggle_traj(True)
    
    target_pose = client.get_pose()
    client.set_named_axis(0, target_pose)
    motion = OnlineMotion(target_pose, vmax=0.4, amax=0.1)

    # 异步运行运动跟踪
    motion_runner = (
        GraphRunner()
        .add_node(0, motion.get_pose)
        .add_node(1, client.set_cartesian, -1)
        .run_async(frequency=50)
    )
    
    def move(x, y, z):
        target_pose.translate(x, y, z, local=False)
        motion.add_pose(target_pose)
        client.set_named_axis(0, target_pose)
    
    with tb.window("control", None, 10, size=(500, 500), pos=(400, 0), frameless=False):
        tb.add_text("text", text="Sim Control")
        tb.add_stretch(1)
        
        tb.add_key_binding(tb.Key.Key_D, lambda: move(0.01, 0, 0))
        tb.add_key_binding(tb.Key.Key_A, lambda: move(-0.01, 0, 0))
        tb.add_key_binding(tb.Key.Key_W, lambda: move(0, 0.01, 0))
        tb.add_key_binding(tb.Key.Key_S, lambda: move(0, -0.01, 0))
        tb.add_key_binding(tb.Key.Key_Q, lambda: move(0, 0, 0.01))
        tb.add_key_binding(tb.Key.Key_E, lambda: move(0, 0, -0.01))

    tb.exec()